import numpy as np
import tensorflow as tf
from keras import backend as K
from keras.models import Sequential,load_model,model_from_json
from keras.layers.core import Dense, Activation, Dropout
from keras.layers import Convolution2D, MaxPooling2D, Flatten
from keras.layers.normalization import BatchNormalization
from keras.optimizers import RMSprop, SGD
from keras.datasets import mnist
from keras.utils import np_utils
from rbflayer import RBFLayer, InitCentersRandom
from initializer import InitFromFile


def accuracy_score(y1, y2):
    assert y1.shape == y2.shape

    y1_argmax = np.argmax(y1, axis=1)
    y2_argmax = np.argmax(y2, axis=1)
    score = sum(y1_argmax == y2_argmax)
    return (score / len(y1)) * 100


def add_rbf_layer(model, betas, X_train, Y_train, X_test, Y_test):
    """ Create a new model as a copy of model + RBF network. Train
    it on [X_train, Y_train], reports its test accuracy and returns
    the new model.
    """

    newmodel = Sequential()
    for i in range(len(model.layers)):  # copy原模型
        newmodel.add(model.layers[i])

    #    for layer in newmodel.layers:
    #        layer.trainable = False

    rbflayer = RBFLayer(300, betas=betas)
    newmodel.add(rbflayer)
    newmodel.add(Dense(10, use_bias=False, name="dense_rbf"))
    newmodel.add(Activation('softmax', name="Activation_rbf"))

    newmodel.compile(loss='categorical_crossentropy',  # 使用compile方法编译模型
                     optimizer=RMSprop(),
                     metrics=['acc'])

    newmodel.summary()

    # model.compile(loss='mean_squared_error',
    #              optimizer=SGD(lr=0.1, decay=1e-6))

    newmodel.fit(X_train, Y_train,  # 训练模型
                 batch_size=128,
                 epochs=3,
                 verbose=1)

    Y_pred = newmodel.predict(X_test)  # 使用训练好的模型进行预测
    print("Test Accuracy: ", accuracy_score(Y_pred, Y_test))

    return newmodel

def load_data():
    """Loads the MNIST dataset.

    # Arguments
        path: path where to cache the dataset locally
            (relative to ~/.keras/datasets).

    # Returns
        Tuple of Numpy arrays: `(x_train, y_train), (x_test, y_test)`.
    """
    path  = './models/mnist.npz'
    f = np.load(path)
    x_train, y_train = f['x_train'], f['y_train']
    x_test, y_test = f['x_test'], f['y_test']
    f.close()
    return (x_train, y_train), (x_test, y_test)

if __name__ == "__main__":

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    K.set_session(sess)

    # load and transform mnist data

    (X_train, y_train), (X_test, y_test) = load_data()


    X_train = X_train.reshape(60000, 28, 28, 1)

    X_train = X_train.astype('float32')
    X_train /= 255

    X_test = X_test.reshape(10000, 28, 28, 1)

    X_test = X_test.astype('float32')
    X_test /= 255

    Y_train = np_utils.to_categorical(y_train, 10)
    Y_test = np_utils.to_categorical(y_test, 10)

    # load model from file
    print("Load model ... ", end="")
    model = model_from_json(open("models/cnn.json").read(),
                            {'InitFromFile': InitFromFile})
    model.build((None, 28, 28, 1))
    model.load_weights("models/cnn_weights.h5")
    # create and learn new model
    newmodel = add_rbf_layer(model, 2.0, X_train, Y_train, X_test, Y_test)

    # save new model to file
    json_string = newmodel.to_json()
    with open("models/cnnrbf.json", "w") as f:
        f.write(json_string)
    newmodel.save_weights("models/cnnrbf_weights.h5")

    print("---test----")
    m = load_model(("./models/cnnrbf_weights.h5"), custom_objects={'RBFLayer': RBFLayer})

    Y_pred = m.predict(X_test)
    print("Test Accuracy: ", accuracy_score(Y_pred, Y_test))
